#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
【节点名称】：
    CameraCalibrationNode
【依赖项安装】：
    pip install spirems
    git clone -b spirecv2-dds https://gitee.com/amovlab/SpireCV.git
    cd SpireCV && pip install -e .
【订阅类型】：
    sensor_msgs::CompressedImage （输入图像）
【发布类型】：
    暂无数据发布
  
【构造参数说明】：
    parameter_file (str): 全局参数文件
    sms_shutdown (bool): 是否接收全局关闭信号,如果需要长期后台执行,建议设置为False
    specified_input_topic (str): 指定输入的话题地址
    realtime_det (bool): 是否是实时,设置为True可以降低延迟,但可能会产生丢帧
【节点参数】：
    square_size : 图像标定板的方块真实宽度, 默认 10 (mm)
    output_file : 标定结果文件名称，默认 calib.yaml
【备注】：
    无
"""


import numpy as np
import cv2 as cv
# local modules
from common import splitfn
# built-in modules
import yaml 
import threading
from spirems import Subscriber, sms2cvimg, Logger
from spirecv.base.BaseNode import BaseNode
from spirecv.dataloader.CameraNode import CameraNode
from queue import Queue
from typing import Union
import json
import sys
import argparse
import time


obj_points   = []
img_points   = []



def save_camera_params_to_json(filename, imageSize, flags, cameraMatrix, distCoeffs):
    # 获取当前时间并格式化为字符串
    calibration_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    cameraMatrix_f = cameraMatrix.flatten()
    distCoeffs_f = distCoeffs.flatten()

    # 这里的相机外参默认全为0
    R = np.identity(3)
    t = np.zeros((3, 1))
    extrinsicMatrix = np.hstack((R, t))
    projectionMatrix = cameraMatrix @ extrinsicMatrix
    projectionMatrix_f = projectionMatrix.flatten()

    # 准备一个字典来存储相机参数
    camera_params ={}
    camera_params["/CameraNode/calibration_time"] = calibration_time
    camera_params["/CameraNode/image_width"] = imageSize[0]
    camera_params["/CameraNode/image_height"] = imageSize[1]
    camera_params["/CameraNode/flags"] = flags
    camera_params["/CameraNode/flags_description"] =  _flags_to_description(flags)
    camera_params["/CameraNode/camera_matrix"] =  cameraMatrix_f.tolist()
    camera_params["/CameraNode/distortion_coefficients"] = distCoeffs_f.tolist()
    camera_params["/CameraNode/projection"] = projectionMatrix_f.tolist()

    # 写入 JSON 文件
    with open(filename, 'w') as file:
        json.dump(camera_params, file, indent=4)


def _flags_to_description(flags):
    # 将标志位转换为描述字符串
    flag_strings = []
    if flags & cv.CALIB_USE_INTRINSIC_GUESS:
        flag_strings.append("use_intrinsic_guess")
    if flags & cv.CALIB_FIX_ASPECT_RATIO:
        flag_strings.append("fix_aspect_ratio")
    if flags & cv.CALIB_FIX_PRINCIPAL_POINT:
        flag_strings.append("fix_principal_point")
    if flags & cv.CALIB_ZERO_TANGENT_DIST:
        flag_strings.append("zero_tangent_dist")
    return ", ".join(flag_strings)


def Calibrate(img, square_size, imw, imh, output_file):
    
    square_size = float(square_size)
    pattern_size = (9, 6)
    pattern_points = np.zeros((np.prod(pattern_size), 3), np.float32)
    pattern_points[:, :2] = np.indices(pattern_size).T.reshape(-1, 2)
    pattern_points *= square_size
    output_file = str(output_file)
    
    while True:
        img = cv.resize(img, (imw, imh)) 
        frame = img
        
        cv.putText(frame, "Press 'c' to add current frame. 'ESC' to finish and calibrate", (10, 20), cv.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2, lineType=cv.LINE_AA)  
        found, corners = cv.findChessboardCorners(img, pattern_size)
        if found:
            img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
            term = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_COUNT, 30, 0.1)
            cv.cornerSubPix(img, corners, (5, 5), (-1, -1), term)
            cv.drawChessboardCorners(frame, pattern_size, corners, found)
            
        key = cv.waitKey(1) & 0xFF  
  
        if key == ord('c'):  
            # 将当前帧添加到截图列表中
            if corners is not None:
                print("Frame captured.")
                img_points.append(corners.reshape(-1, 2))
                obj_points.append(pattern_points) 
                
        # 检查是否按下退出键（例如 'ESC'）  
        elif key == 27:
            if len(img_points) < 4:
                print("Not enough captures for calibration")
            else:
                # calculate camera distortion
                rms, camera_matrix, dist_coefs, _rvecs, _tvecs = cv.calibrateCamera(obj_points, img_points, (imw, imh), None, None)
            
                flags = cv.CALIB_FIX_ASPECT_RATIO | cv.CALIB_USE_INTRINSIC_GUESS
                save_camera_params_to_json(output_file, (imw, imh), flags, camera_matrix, dist_coefs)
    
                print("\nRMS:", rms)
                print("camera matrix:\n", camera_matrix)
                print("distortion coefficients: ", dist_coefs.ravel())
                print("Calibration saved to : ", output_file)
        
                print('Done!')
            break

            
        
        return frame


class CameraCalibrationNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        specified_input_topic: str = None,
        realtime_det: bool = True
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(self, self.__class__.__name__, job_name, ip=ip, port=port, param_dict_or_file=param_dict_or_file, sms_shutdown=sms_shutdown)
        input_url = '/' + job_name + '/sensor/image_raw'
        if specified_input_topic is not None:
            input_url = specified_input_topic 

        self._image_reader = Subscriber(
            input_url, 'sensor_msgs::CompressedImage', self.image_callback,
            ip=ip, port=port
        )
        
        self.realtime_det = realtime_det
        self.job_queue = Queue()
        self.queue_pool.append(self.job_queue)
        self.logger = Logger(self.__class__.__name__)
        
        self.square_size = self.get_param('square_size', 10)
        self.output_file = self.get_param('output_file', 'calib.json')
        self.logger.info("square_size: " + str(self.square_size))
        self.logger.info("output_file: " + self.output_file)

        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()
        self.logger.quit()

    def image_callback(self, msg):
        if self.realtime_det:
            if not self.job_queue.empty():
                self.job_queue.queue.clear()
        self.job_queue.put(msg)

    def run(self):
        while self.is_running():
            msg = self.job_queue.get(block=True)
            if msg is None:
                break
            
            img = sms2cvimg(msg)
            frame = Calibrate(img, self.square_size, img.shape[1], img.shape[0], self.output_file)
                
            if frame is None:  
                # print(f"Error: Unable to read image.")  
                break  # 终止程序，并返回错误码 1  
            else:
                cv.imshow('out', frame)

        cv.destroyAllWindows() 
        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'config',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--input-topic',
        type=str,
        default=None,
        help='Specified SpireCV2 Input Image Topic')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    args = parser.parse_args()
    print("config:", args.config)
    print("job-name:", args.job_name)
    print("specified-input-topic:", args.input_topic)

    algorithm = CameraCalibrationNode(args.job_name, specified_input_topic=args.input_topic, param_dict_or_file=args.config, ip=args.ip, port=args.port)
    

    
    
    
    
